use dir_writer::{FileCollector, GeneratorArgs, IntermediateRepr, LanguageFeatures};
use functions::{render_client, render_runtime};
use generated_types::render_rb_types;

use crate::{
    functions::render_globals,
    generated_types::{
        render_rb_stream_types_utils, render_rb_type_builder, render_rb_types_utils,
    },
};

mod functions;
mod generated_types;
mod ir_to_rb;
mod package;
mod r#type;
mod utils;

#[derive(Default)]
pub struct RbLanguageFeatures {
    requires: std::sync::Mutex<std::collections::HashMap<std::path::PathBuf, Vec<String>>>,
    requires_relative: std::sync::Mutex<std::collections::HashMap<std::path::PathBuf, Vec<String>>>,
}

impl RbLanguageFeatures {
    fn add_import(&self, path: &str, import: &str, relative: bool) {
        let mut map = if relative {
            self.requires_relative.lock().unwrap()
        } else {
            self.requires.lock().unwrap()
        };
        map.entry(std::path::Path::new(path).to_path_buf())
            .or_insert_with(Vec::new)
            .push(import.to_string());
    }
}

impl LanguageFeatures for RbLanguageFeatures {
    const CONTENT_PREFIX: &'static str = r#"
# typed: strict
# ----------------------------------------------------------------------------
#
#  Welcome to Baml! To use this generated code, please run the following:
#
#  $ gem install baml
#
# ----------------------------------------------------------------------------

# This file was generated by BAML: please do not edit it. Instead, edit the
# BAML files and re-generate this code using: baml-cli generate
# baml-cli is available with the baml package.
# gem install baml
        "#;

    fn name() -> &'static str {
        "ruby/sorbet"
    }

    fn on_file_created(
        &self,
        _path: &std::path::Path,
        _content: &mut String,
    ) -> anyhow::Result<()> {
        // Do nothing we'll do this in on_file_finished
        self.add_import(_path.to_str().unwrap(), "sorbet-runtime", false);
        self.add_import(_path.to_str().unwrap(), "baml", false);
        Ok(())
    }

    fn on_file_finished(&self, path: &std::path::Path, content: &mut String) -> anyhow::Result<()> {
        *content = {
            let mut new_content = self.content_prefix().to_string();
            if let Some(requires) = self.requires.lock().unwrap().get(path) {
                new_content.push('\n');
                for require in requires {
                    new_content.push_str(&format!("require \"{require}\"\n"));
                }
            }
            if let Some(requires) = self.requires_relative.lock().unwrap().get(path) {
                new_content.push('\n');
                for require in requires {
                    new_content.push_str(&format!("require_relative \"{require}\"\n"));
                }
            }
            new_content.push('\n');
            new_content.push_str("module BamlClient\n");
            for line in content.split("\n") {
                if line.trim().is_empty() {
                    new_content.push('\n');
                } else {
                    new_content.push_str(&format!("  {line}\n"));
                }
            }
            new_content.push_str("\nend\n");
            new_content
        };
        Ok(())
    }

    fn generate_sdk_files(
        &self,
        collector: &mut FileCollector<Self>,
        ir: std::sync::Arc<IntermediateRepr>,
        args: &GeneratorArgs,
    ) -> Result<(), anyhow::Error> {
        let pkg = package::CurrentRenderPackage::new("BamlClient", ir.clone());
        let _file_map = args.file_map_as_json_string()?;

        // collector.add_file("b.rb", render_init(&pkg, &args.default_client_mode)?)?;
        // collector.add_file("inlinedbaml.rb", render_source_files(file_map)?)?;
        collector.add_file("runtime.rb", render_runtime(&pkg)?)?;
        self.add_import("runtime.rb", "type_builder", true);
        // collector.add_file("tracing.rb", render_tracing(&pkg)?)?;
        collector.add_file("globals.rb", render_globals(&pkg)?)?;
        // collector.add_file("config.rb", render_config(&pkg)?)?;
        let functions = ir
            .functions
            .iter()
            .map(|f| ir_to_rb::functions::ir_function_to_rb(f, &pkg))
            .collect::<Vec<_>>();
        collector.add_file("client.rb", render_client(&functions, &pkg)?)?;
        // collector.add_file("parser.rb", render_parser(&functions, &pkg)?)?;

        let rb_classes = ir
            .walk_classes()
            .map(|c| ir_to_rb::classes::ir_class_to_rb(c.item, &pkg))
            .collect::<Vec<_>>();
        let enums = ir
            .walk_enums()
            .map(|e| ir_to_rb::enums::ir_enum_to_rb(e.item, &pkg))
            .collect::<Vec<_>>();
        let type_aliases = ir.walk_type_aliases().collect::<Vec<_>>();
        let mut rb_type_aliases = type_aliases
            .iter()
            .map(|c| ir_to_rb::type_aliases::ir_type_alias_to_rb(c.item, &pkg))
            .collect::<Vec<_>>();
        rb_type_aliases.sort_by(|a, b| a.name.cmp(&b.name));

        // pkg.set("baml_client.type_map");
        // collector.add_file("type_map.rb", render_type_map(&rb_classes, &enums)?)?;

        pkg.set("BamlClient");
        collector.add_file(
            "type_builder.rb",
            render_rb_type_builder(&rb_classes, &enums)?,
        )?;

        pkg.set("BamlClient.Types");
        collector.add_file("types.rb", "module Types\n")?;
        collector.append_to_file("types.rb", &render_rb_types_utils(&pkg)?)?;
        collector.append_to_file("types.rb", &render_rb_types(&enums, &pkg)?)?;
        collector.append_to_file("types.rb", &render_rb_types(&rb_classes, &pkg)?)?;
        collector.append_to_file("types.rb", &render_rb_types(&rb_type_aliases, &pkg)?)?;
        collector.append_to_file("types.rb", "\nend\n")?;

        let mut rb_stream_type_aliases = type_aliases
            .iter()
            .map(|c| ir_to_rb::type_aliases::ir_type_alias_to_rb_stream(c.item, &pkg))
            .collect::<Vec<_>>();
        rb_stream_type_aliases.sort_by(|a, b| a.name.cmp(&b.name));

        let rb_classes = ir
            .walk_classes()
            .map(|c| ir_to_rb::classes::ir_class_to_rb_stream(c.item, &pkg))
            .collect::<Vec<_>>();

        pkg.set("BamlClient.StreamTypes");
        collector.add_file("stream_types.rb", "module StreamTypes\n")?;
        collector.append_to_file("stream_types.rb", &render_rb_stream_types_utils(&pkg)?)?;
        collector.append_to_file("stream_types.rb", &render_rb_types(&rb_classes, &pkg)?)?;
        collector.append_to_file(
            "stream_types.rb",
            &render_rb_types(&rb_stream_type_aliases, &pkg)?,
        )?;
        collector.append_to_file("stream_types.rb", "\nend\n")?;

        Ok(())
    }
}

// DISABLED: Ruby tests are not working yet
// #[cfg(test)]
// mod ruby_tests {
//     use test_harness::{create_code_gen_test_suites, TestLanguageFeatures};

//     impl TestLanguageFeatures for crate::RbLanguageFeatures {
//         fn test_name() -> &'static str {
//             "ruby"
//         }
//     }

//     create_code_gen_test_suites!(crate::RbLanguageFeatures);
// }

#[cfg(test)]
mod tests {
    #[test]
    fn test_name() {
        use std::str::FromStr;

        use dir_writer::LanguageFeatures;

        let gen_type = baml_types::GeneratorOutputType::from_str(crate::RbLanguageFeatures::name())
            .expect("RbLanguageFeatures name should be a valid GeneratorOutputType");
        assert_eq!(gen_type, baml_types::GeneratorOutputType::RubySorbet);
    }
}
